//
//  CPABEInterface.m
//  iHealthEHR
//
//  Created by Joseph Ayo Akinyele on 4/21/10.
//  Copyright 2010 __MyCompanyName__. All rights reserved.
//

#import "CPABEInterface.h"

#define SESSION_KEY_LEN	16
// #define SIZE_MAX 4096
#define SIZE_MED 512

void report_error(char* action, FENC_ERROR result)
{
	printf("%s...\n\t%s (%d)\n", action, libfenc_error_to_string(result), result);
}

void print_buffer_as_hex(uint8* data, size_t len)
{
	int i;
	
	for (i = 0; i < len; i++) {
		printf("%02x ", data[i]);
	}
	printf("\n");
}


ssize_t read_file(FILE *f, char** out) {
	
	if(f != NULL) {
		/* See how big the file is */
		fseek(f, 0L, SEEK_END);
		ssize_t out_len = ftell(f);
		// printf("out_len: %zd\n", out_len);
		if(out_len <= SIZE_MAX) {
			/* allocate that amount of memory only */
			if((*out = (char *) malloc(out_len+1)) != NULL) {
				fseek(f, 0L, SEEK_SET);
				fread(*out, sizeof(char), out_len, f);
				return out_len;
			}
		}
	}
	return 0;
}

/* Should return an AES session key */
unsigned char * _ABEdecryption(fenc_context *context, fenc_key *secret_key, const char *abe_cipher_64, size_t *keyLength)
{
	FENC_ERROR result;
	size_t abeLength;
	char *cipher_base64;
	fenc_ciphertext ciphertext;
	fenc_plaintext aes_session_key;
	char session_key[SESSION_KEY_LEN];
	
	/* clear structures */
	memset(&ciphertext, 0, sizeof(fenc_ciphertext));
	memset(&aes_session_key, 0, sizeof(fenc_plaintext));
	memset(session_key, 0, SESSION_KEY_LEN);
	
	cipher_base64 = NewBase64Decode((const char *) abe_cipher_64, strlen(abe_cipher_64), &abeLength);
	ciphertext.data = (uint8 *) cipher_base64;
	ciphertext.data_len = abeLength;
	ciphertext.max_len = abeLength;
	
	/* Descrypt the resulting ciphertext. */
	// ciphertext = input
	result = libfenc_decrypt(context, &ciphertext, secret_key, &aes_session_key);
	if (result == FENC_ERROR_NONE) {
		/* make sure key is random...? */
		if (memcmp(aes_session_key.data, session_key, aes_session_key.data_len) != 0) {
			result = FENC_ERROR_UNKNOWN;
		}
	}
	report_error("Decrypting the ciphertext", result);
	
	printf("Decrypted session key is: ");
	print_buffer_as_hex(aes_session_key.data, aes_session_key.data_len);
	
	uint8 *thekey;
	if((thekey = malloc(aes_session_key.data_len) ) == NULL) {
		return NULL;
	}
	
	memcpy(thekey, aes_session_key.data, aes_session_key.data_len);
	/* caller must free */
	*keyLength = aes_session_key.data_len;
	/* free others */
	free(cipher_base64);
	return thekey;
}

char* _RetrieveABEPolicy(fenc_context *context, const char *abe_cipher_64)
{
	FENC_ERROR result;
	char *cipher_base64;
	size_t abeLength;
	fenc_ciphertext ciphertext;
	uint8 *policy = SAFE_MALLOC(SIZE_MED);
	
	/* clear structures */
	memset(&ciphertext, 0, sizeof(fenc_ciphertext));
	memset(policy, 0, SIZE_MED);
	
	cipher_base64 = NewBase64Decode((const char *) abe_cipher_64, strlen(abe_cipher_64), &abeLength);
	ciphertext.data = (uint8 *) cipher_base64;
	ciphertext.data_len = abeLength;
	ciphertext.max_len = abeLength;
	
	result = libfenc_retrieve_attribute_policy(context, &ciphertext, policy, SIZE_MED);
	report_error("Retrieve policy string from ciphertext", result);
	
	printf("Policy: '%s'\n", policy);
	
	/* caller must free */
	return (char *) policy;
}


char* _AESdecryption(uint8 *session_key, const char *aes_cipher)
{
	/* decode the aesblob64 */
	size_t aesLength;
	char *aesblob = NewBase64Decode(aes_cipher, strlen(aes_cipher), &aesLength);
	
	/* use the PSK to encrypt using openssl functions here */
	AES_KEY skey;
	uint8 iv[SESSION_KEY_LEN*4];
	char aes_result[aesLength+1];
	AES_set_decrypt_key(session_key, 8*SESSION_KEY_LEN, &skey);
	memset(iv, 0, SESSION_KEY_LEN*4);
	memset(aes_result, 0, aesLength+1);
	AES_cbc_encrypt((uint8 *) aesblob, (uint8 *) aes_result, aesLength, &skey, iv, AES_DECRYPT);
	/* base-64 both ciphertext and write to the stdout -- in XML? */
	
	// printf("Plaintext: %s\nSize: %zd\n", aes_result, aesLength);
	char *plain;
	if((plain = malloc(aesLength)) == NULL) 
		return NULL;
	
	memcpy(plain, aes_result, aesLength);
	free(aesblob);
	return plain;	
}

#pragma mark -
#pragma mark Interface init methods

@implementation CPABEInterface

-(void) Initialize
{
	FENC_ERROR result;
	pairing_t pairing;
	FILE *fp;
	/* Clear data structures. */
	memset(&context, 0, sizeof(fenc_context));
	memset(&group_params, 0, sizeof(fenc_group_params));
	memset(&global_params, 0, sizeof(fenc_global_params));	
	
	/* Initialize the library. */
	result = libfenc_init();
	/* Create a Sahai-Waters context. */
	result = libfenc_create_context(&context, FENC_SCHEME_WATERSCP);	
	report_error("Initalize Waters CP scheme", result);
	
	/* Load group parameters from a file. */
	NSString *path = [[NSBundle mainBundle] pathForResource:@"d224" ofType:@"param"];
	fp = fopen([path cStringUsingEncoding:NSASCIIStringEncoding], "r");
	if (fp != NULL) {
		libfenc_load_group_params_from_file(&group_params, fp);
		libfenc_get_pbc_pairing(&group_params, pairing);
	} else {
		perror("Could not open parameters file.\n");
		return;
	}
	fclose(fp);
	
	/* Set up the global parameters. */
	result = context.generate_global_params(&global_params, &group_params);
	report_error("Loading global parameters", result);
	
	result = libfenc_gen_params(&context, &global_params);
	report_error("Generating scheme parameters and secret key", result);
}

/* Reads the public parameters and secret parameters from the filesystem (temporary)
 */
-(BOOL) ReadParameters
{
	FILE *fp;
	FENC_ERROR result;
	char *public_buf = NULL; // , *m_secret_buf = NULL;
	ssize_t public_len; // , m_secret_len;
	size_t serialized_len;
	NSString *path = [[NSBundle mainBundle] pathForResource:@"public" ofType:@"param"];
	fp = fopen([path cStringUsingEncoding:NSASCIIStringEncoding], "r");
	if(fp != NULL) {
		if((public_len = read_file(fp, &public_buf)) > 0) {
			// printf("Public parameters: %s\n", public_buf);
			uint8 *bin_public_buf = NewBase64Decode((const char *) public_buf, public_len, &serialized_len);
			
			/* Import the parameters from binary buffer: */
			result = libfenc_import_public_params(&context, bin_public_buf, serialized_len);
			report_error("Importing public parameters", result);
			
			free(bin_public_buf);
			free(public_buf);
		}
	}
	else {
		perror("Could not open public parameters.\n");
		return NO;
	}	
	fclose(fp);
	/* temporary - read in master secret parameters 
	NSString *path1 = [[NSBundle mainBundle] pathForResource:@"master_secret" ofType:@"param"];
	fp = fopen([path1 cStringUsingEncoding:NSASCIIStringEncoding], "r");
	if(fp != NULL) {
		if((m_secret_len = read_file(fp, &m_secret_buf)) > 0) {
			// printf("Master secret parameters: %s\n", m_secret_buf);
			uint8 *bin_secret_buf = NewBase64Decode((const char *) m_secret_buf, m_secret_len, &serialized_len);
			
			// Import the parameters from binary buffer 
			result = libfenc_import_secret_params(&context, bin_secret_buf, serialized_len, NULL, 0);
			report_error("Importing secret parameters", result);
			free(bin_secret_buf);
			free(m_secret_buf);
		}
	}
	else {
		perror("Could not open public parameters.\n");
		return FALSE;
	}	
	fclose(fp); */
	
	return YES;
}

-(void) CleanUp
{
	FENC_ERROR result;
	/* Destroy the context. */
	result = libfenc_destroy_context(&context);
	report_error("Destroying the context", result);	
	
	/* Shutdown the library. */
	result = libfenc_shutdown();
	report_error("Shutting down library", result);
}

#pragma mark -
#pragma mark key-generator and decryption methods
/* Note: this method should be retrieving the user's private key from the
iPhone keychain. Instead, we retrieve from the local filesystem. THIS IS ONLY A TEST!
 */
#define SIZE 250
-(BOOL) RetrieveKey:(NSString *) password with:(NSString *) username
{
	FENC_ERROR result;
	FILE *fp;
	char *u_secret_buf = NULL;
	ssize_t u_secret_len;
	size_t serialized_len;
	/* read file */
	NSString *path1 = [[NSBundle mainBundle] pathForResource:@"user_secret" ofType:@"key"];
	fp = fopen([path1 cStringUsingEncoding:NSASCIIStringEncoding], "r");
	if(fp != NULL) {
		if((u_secret_len = read_file(fp, &u_secret_buf)) > 0) {
			uint8 *bin_secret_buf = NewBase64Decode((const char *) u_secret_buf, u_secret_len, &serialized_len);
			
			//printf("User's base-64 decoded secret key:\n");
			//print_buffer_as_hex(bin_secret_buf, serialized_len);
			
			/* Import the parameters from binary buffer */
			result = libfenc_import_secret_key(&context, &secret_key, bin_secret_buf, serialized_len);
			report_error("Importing secret key", result);
			free(bin_secret_buf);
			free(u_secret_buf);
		}
	}
	else {
		perror("Could not open user_secret key.\n");
		return NO;
	}	
	fclose(fp);
	
	return YES;
}

-(void) GenerateStaticKey
{
	FENC_ERROR result;
	uint8 output[SIZE];
	size_t output_len;
	memset(&func_list_input, 0, sizeof(fenc_function_input));
	char *attr[9] = {"ONE", "TWO", "THREE", "FOUR", "FIVE", "SIX", "SEVEN", "EIGHT"};
	libfenc_create_attribute_list_from_strings(&func_list_input, attr, 8);

	fenc_attribute_list_to_buffer((fenc_attribute_list*)(func_list_input.scheme_input), output, SIZE, &output_len);
	printf("Attribute list: %s\n", output);
	
	result = libfenc_extract_key(&context, &func_list_input, &secret_key);
	report_error("Retrieving user's decryption key", result);
}

/* Note: caller must free the returned object */
-(NSString *) Decrypt:(NSString *) abe_cipher and:(NSString *) aes_cipher
{	
	uint8 *session_key;
	char *plain;
	size_t key_len;
	session_key = _ABEdecryption(&context, &secret_key, [abe_cipher cStringUsingEncoding:NSASCIIStringEncoding], &key_len);
	printf("AES Session key: \t");
	print_buffer_as_hex(session_key, key_len);
	
	plain = _AESdecryption(session_key, [aes_cipher cStringUsingEncoding:NSASCIIStringEncoding]);
	// printf("THE MESSAGE: '%s'\n", plain);

	NSString *plaintext = [[NSString alloc] initWithFormat:@"%s", plain];
	free(session_key);
	free(plain);
	return plaintext;
}

/* Note: caller must free the returned object */
-(NSString *) ReturnPolicy:(NSString *) abe_cipher
{
	char *policy;
	policy = _RetrieveABEPolicy(&context, [abe_cipher cStringUsingEncoding:NSASCIIStringEncoding]);
	
	NSString *policy_str = [[NSString alloc] initWithFormat:@"%s", policy];
	free(policy);
	
	return policy_str;
}

@end
